import sensor, image, time, math
import pyb
from pyb import UART, LED
import ulab as np
import gc

# LED状态指示初始化
pyb.LED(1).on()   # 红色LED指示电源
pyb.LED(2).off()  # 绿色LED指示摄像头状态
pyb.LED(3).off()  # 蓝色LED指示通信状态

# #################### 常量定义 ####################
# 系统操作模式字典（状态机基础）
OPERATION_MODE = {
    'STANDBY': 0,                 # 待机模式
    'LINE_TRACKING': 1,           # 巡线模式
    'COLOR_DETECTION': 2,         # 物料颜色识别
    'COARSE_AREA_DETECTION': 3,   # 粗加工区定位
    'BUFFER_AREA_DETECTION': 4,   # 缓存区定位
    'QR_CODE_SCAN': 5             # 二维码扫描
}

# UART通信常量
UART_HEADER = 0xAA
UART_FOOTER = 0xBB

# 命令映射表 (主命令)
COMMAND_MAP = {
    0x01: 'LINE_TRACKING',
    0x02: 'COLOR_DETECTION',
    0x03: 'COARSE_AREA_DETECTION',
    0x04: 'BUFFER_AREA_DETECTION',
    0x05: 'QR_CODE_SCAN',
    0x00: 'STANDBY'
}

# 颜色子命令映射 (用于颜色检测模式)
COLOR_COMMAND_MAP = {
    0x01: 'red',
    0x02: 'green',
    0x03: 'blue'
}

# 优化后的颜色阈值 - LAB色彩空间
COLOR_THRESHOLDS = {
    'red': [(28, 71, 5, 63, -8, 51)],          # 优化红色阈值
    #'green': [(64, 75, -63, -20, 8, 50)],       # 优化绿色阈值
    'green': [(52, 67, -63, -24, -15, 53)],
    #'blue': [(30, 67, -3, 10, -54, -25)],      # 统一蓝色阈值
    'blue': [(23, 78, 7, 31, -62, -45)],
    'line_tracking': [(0, 30, 0, 64, -128, -20)],  # 黑线阈值
    # 'coarse_red': [(17, 56, 3, 13, -3, 21)],
    # 'coarse_green': [(23, 57, -36, -19, -9, 19)],
    # 'coarse_blue': [(22, 62, -22, -11, -19, -1)]
}



# 颜色识别结果代码 (与命令代码区分开)
COLOR_RESULT_CODES = {
    'red': 0x1,
    'green': 0x2,
    'blue': 0x3
}

DETECTION_PARAMS = {
    'line_regression': {
        'x_stride': 2,
        'y_stride': 2,
        'threshold': 1000,
        'theta_margin': 25,
        'rho_margin': 25
    },
    'coarse_area_circle': {
        'r_min': 30,
        'r_max': 95,
        'samples_needed': 6,  # 增加采样点
        'area_threshold': 3000  # 降低阈值，增加识别率
    }
}

# #################### 初始化设置 ####################
# 摄像头初始化
sensor.reset()
sensor.set_pixformat(sensor.RGB565)
sensor.set_framesize(sensor.QQVGA)    # 160x120 分辨率 (更正注释)
sensor.skip_frames(time=2000)
sensor.set_auto_gain(False)
sensor.set_auto_whitebal(False)
sensor.set_auto_exposure(False, exposure_us=8000)  # 适当增加曝光时间
sensor.set_contrast(1)
# sensor.set_brightness(0)             # 设置亮度
# sensor.set_saturation(1)             # 增加饱和度以便更好地区分颜色

# UART初始化
uart = UART(3, 115200)
uart.init(115200, bits=8, parity=None, stop=1)

# 全局变量
current_operation = OPERATION_MODE['STANDBY']
target_color = None  # 目标检测颜色
line_x_samples = [0] * 15
line_y_samples = [0] * 15
sample_counter = 0
# DEBUG_MODE = True  # 调试模式开关
# DEBUG_TARGET_COLOR = 'blue'

# *** 新增：用于控制颜色检测发送间隔的变量 ***
last_color_send_time = 0
SEND_INTERVAL_MS = 1000 # 发送间隔，100ms = 0.1s

# 圆检测参数 (需要根据实际目标和距离调整)
CIRCLE_THRESHOLD = 2500   # 圆检测的边缘数量阈值，越高越严格
CIRCLE_R_MIN = 20         # 检测圆的最小半径 (像素)
CIRCLE_R_MAX = 80         # 检测圆的最大半径 (像素)
CIRCLE_MAG_THRESHOLD = 6  # 边缘梯度强度阈值，越高越不容易受噪声影响
CIRCLE_MERGE_DIST = 10    # 合并距离相似的圆

# EMA 平滑参数
EMA_ALPHA = 0.3  # 平滑因子 (0 < alpha <= 1)，越小越平滑，但响应越慢
smoothed_x = None
smoothed_y = None
smoothed_r = None
target_found = False # 标记是否找到目标

# #################### 功能函数 ####################
def send_data(*payload):
    """发送数据包到STM32（优化版）"""
    # 仅清空发送缓冲区，保留接收数据
    uart.writechar(UART_HEADER)
    for byte in payload:
        uart.writechar(byte & 0xFF)
    uart.writechar(UART_FOOTER)
    pyb.delay(2) # 短暂延时确保发送完成
def parse_command():
    """解析从STM32接收的命令（增强版）"""
    if uart.any():
        raw_data = uart.read()
        data_bytes = bytearray(raw_data)

        # 调试输出原始数据
        print("Raw RX:", [hex(b) for b in data_bytes])

        # 查找有效数据包起始位置
        start_idx = -1
        for i in range(len(data_bytes)):
            if data_bytes[i] == UART_HEADER:
                start_idx = i
                break

        # 未找到帧头则直接返回
        if start_idx == -1:
            return {'valid': False}

        # 截取有效数据段
        data_bytes = data_bytes[start_idx:]

        # 查找帧尾位置
        end_idx = -1
        for i in range(len(data_bytes)):
            if data_bytes[i] == UART_FOOTER:
                end_idx = i
                break

        # 未找到完整数据包
        if end_idx == -1 or end_idx <= start_idx:
            return {'valid': False}

        # 提取完整数据包
        packet = data_bytes[:end_idx+1]

        # 校验数据包长度
        if len(packet) < 4:  # 最小有效包：AA CMD DATA BB
            return {'valid': False}

        # 解析命令
        result = {
            'valid': False,
            'command': 'UNKNOWN',
            'color': None
        }

        cmd_byte = packet[1]
        if cmd_byte in COMMAND_MAP:
            result['command'] = COMMAND_MAP[cmd_byte]
            result['valid'] = True

            # 处理颜色子命令
            if result['command'] in ['COLOR_DETECTION', 'COARSE_AREA_DETECTION'] and len(packet) >= 4:
                color_cmd = packet[2]
                if color_cmd in COLOR_COMMAND_MAP:
                    result['color'] = COLOR_COMMAND_MAP[color_cmd]

        return result

    return {'valid': False}




def detect_specific_color(img, target_color):
    """检测指定颜色的物体，返回检测结果"""
    # 定义检测参数
    min_blob_area = 1000      # 最小面积要求 (原值)
    # QQVGA is 160x120, adjust ROI accordingly
    # roi = (20, 20, 120, 80)  # (x, y, w, h) ROI区域, 相对于160x120
    roi = (int(img.width()*0.1), int(img.height()*0.1), int(img.width()*0.8), int(img.height()*0.8)) # 动态ROI

    # 在ROI区域绘制矩形框(调试用)
    img.draw_rectangle(roi, color=(255,255,255), thickness=1)

    # 获取目标颜色的阈值
    threshold = COLOR_THRESHOLDS.get(target_color)
    if not threshold:
        print(f"未知颜色: {target_color}")
        return None

    # 根据颜色选择标记颜色
    if target_color == 'red':
        mark_color = (255, 0, 0)
    elif target_color == 'green':
        mark_color = (0, 255, 0)
    else:  # 蓝色
        mark_color = (0, 0, 255)

    # 查找色块
    blobs = img.find_blobs(
        threshold,
        roi=roi,
        pixels_threshold=500,    # 像素阈值 (原值)
        area_threshold=min_blob_area,
        merge=True,
        margin=10
    )

    # 找到有效的色块
    if blobs:
        # 选择最大的色块
        largest_blob = max(blobs, key=lambda b: b.area())

        # 检查形状是否合理(排除长条或奇怪形状) - 可选
        ratio = largest_blob.w() / largest_blob.h() if largest_blob.h() != 0 else 0
        if 0.5 <= ratio <= 2.0:  # 限制宽高比 (原值)
            # 绘制检测结果
            img.draw_rectangle(largest_blob.rect(), color=mark_color, thickness=2)
            img.draw_cross(largest_blob.cx(), largest_blob.cy(), color=mark_color, thickness=2)
            # img.draw_string(largest_blob.cx()-10, largest_blob.cy()-30, target_color.upper(),
            #                color=mark_color, scale=2) # 字符串绘制比较耗时，调试时可选

            # 计算归一化坐标(0-255范围)
            norm_x = int(largest_blob.cx() * 255 / img.width())
            norm_y = int(largest_blob.cy() * 255 / img.height())

            # 返回检测结果
            return {
                'color': target_color,
                'code': COLOR_RESULT_CODES[target_color],  # 使用结果代码，而非命令代码
                'x': norm_x,
                'y': norm_y,
                'area': largest_blob.area()
            }

    return None # 如果没有找到或者形状不合理，返回None

# #################### 主循环 ####################
clock = time.clock()

# 主循环
while True:
    clock.tick()
    img = sensor.snapshot()

    # 清理内存
    gc.collect()


    DEBUG_MODE = True  # 调试开关
    if DEBUG_MODE:
        #current_operation = OPERATION_MODE['COLOR_DETECTION']
        current_operation = OPERATION_MODE['COARSE_AREA_DETECTION']

        target_color = 'blue'
        pyb.LED(2).on()

    # 处理STM32命令
    cmd_result = parse_command()

    if cmd_result['valid']:
        command_str = cmd_result['command']
        # 仅当命令解析有效且命令存在于 OPERATION_MODE 中时才切换状态
        if command_str in OPERATION_MODE:
            # 检查状态是否真的改变，或者目标颜色是否改变，避免不必要的打印和LED闪烁
            if OPERATION_MODE[command_str] != current_operation or \
               (command_str in ['COLOR_DETECTION', 'COARSE_AREA_DETECTION'] and cmd_result['color'] != target_color):

                current_operation = OPERATION_MODE[command_str]
                print(f"切换到 {command_str} 模式")

                # 同时处理需要颜色的模式
                if command_str in ['COLOR_DETECTION', 'COARSE_AREA_DETECTION']:
                    target_color = cmd_result['color']
                    if target_color: # 确保颜色子命令有效
                         print(f"目标颜色: {target_color}")
                         pyb.LED(2).on()
                    else:
                         # 如果颜色子命令无效，则可能需要回到 STANDBY 或报错
                         print("错误：需要颜色子命令但未提供或无效")
                         current_operation = OPERATION_MODE['STANDBY'] # 回到待机
                         pyb.LED(2).off()
                else:
                    target_color = None
                    pyb.LED(2).off() # 关闭绿色LED
            # else:
                # print("状态未改变") # 调试用


    # 颜色检测模式且有指定目标颜色
    if current_operation == OPERATION_MODE['COLOR_DETECTION'] and target_color:
        # 执行特定颜色检测
        result = detect_specific_color(img, target_color)

        # *** 修改部分：定时发送 ***
        current_time = time.ticks_ms()
        # 检查时间间隔是否到达
        if time.ticks_diff(current_time, last_color_send_time) >= SEND_INTERVAL_MS:
            # 只有在检测到结果时才发送
            if result:
                # 闪烁蓝色LED表示发送数据
                pyb.LED(3).toggle()

                # 发送颜色识别结果 (颜色结果代码, X坐标, Y坐标)
                send_data(result['code'], result['x'], result['y'])
                print(f"发送: 颜色={result['color']}, X={result['x']}, Y={result['y']}")

                # 更新上次发送时间戳 (只有发送了才更新)
                last_color_send_time = current_time
            # else:
                # 如果当前帧未检测到，则不发送，也不更新时间戳
                # print("未检测到目标颜色，不发送") # 调试用


    # 修改后的COARSE_AREA_DETECTION处理 (保持不变)

    elif current_operation == OPERATION_MODE['COARSE_AREA_DETECTION']:
            if DEBUG_MODE >= 1:
                print("当前状态: 粗加工区检测")

            target_found_in_frame = False # 重置当前帧查找状态
            best_circle = None

            # --- 核心检测逻辑 ---
            # 寻找圆形（可以使用颜色阈值过滤，但 find_circles 本身基于边缘）
            # 为了提高 find_circles 的鲁棒性，可以先对图像进行一些预处理，
            # 例如二值化，但这里直接在原图上找
            detected_circles = img.find_circles(threshold = CIRCLE_THRESHOLD,
                                                r_min = CIRCLE_R_MIN,
                                                r_max = CIRCLE_R_MAX,
                                                magnitude_threshold = CIRCLE_MAG_THRESHOLD,
                                                x_margin=10, y_margin=10, r_margin=CIRCLE_MERGE_DIST)

            if detected_circles:
                # 选择最合适的圆 (例如：选择半径最大的圆)
                # 你可以根据需要修改选择逻辑 (比如最接近图像中心的圆)
                detected_circles.sort(key=lambda c: c.r(), reverse=True)
                best_circle = detected_circles[0]

                # # --- 检查圆内部颜色 (可选，增加可靠性) ---
                # # 取圆心一小块区域判断颜色是否符合预期
                # COLOR_THRESHOLD = [52, 67, -63, -24, -15, 53]
                # center_roi = (best_circle.x() - 5, best_circle.y() - 5, 10, 10)
                # center_stats = img.get_statistics(roi=center_roi)
                # if COLOR_THRESHOLD[0] <= center_stats.l_mean() <= COLOR_THRESHOLD[1] and \
                #    COLOR_THRESHOLD[2] <= center_stats.a_mean() <= COLOR_THRESHOLD[3] and \
                #    COLOR_THRESHOLD[4] <= center_stats.b_mean() <= COLOR_THRESHOLD[5]:
                #     # 颜色匹配成功
                #     target_found_in_frame = True
                # else:
                #     # 虽然找到圆，但颜色不对
                #     if DEBUG_MODE >= 2:
                #         print("检测到圆，但中心颜色不匹配")
                #     best_circle = None # 标记为未找到
                # # 简化处理：只要找到圆就认为找到
                # target_found_in_frame = True

            # --- 更新平滑数据 (EMA) ---
            if target_found_in_frame and best_circle:
                current_x = best_circle.x()
                current_y = best_circle.y()
                current_r = best_circle.r()

                if smoothed_x is None: # 第一次检测到
                    smoothed_x = current_x
                    smoothed_y = current_y
                    smoothed_r = current_r
                else: # 应用 EMA
                    smoothed_x = EMA_ALPHA * current_x + (1.0 - EMA_ALPHA) * smoothed_x
                    smoothed_y = EMA_ALPHA * current_y + (1.0 - EMA_ALPHA) * smoothed_y
                    smoothed_r = EMA_ALPHA * current_r + (1.0 - EMA_ALPHA) * smoothed_r

                target_found = True # 更新全局目标状态


                if DEBUG_MODE >= 2:
                    print(f"检测到圆: ({current_x}, {current_y}, R={current_r})")
                    print(f"平滑后圆: ({int(smoothed_x)}, {int(smoothed_y)}, R={int(smoothed_r)})")

            else:
                # 当前帧未找到圆
                target_found = False # 如果希望目标状态在丢失后保持一小段时间，需要更复杂的逻辑

                # smoothed_x, smoothed_y, smoothed_r 保持不变，代表最后一次的平滑位置
                # 如果长时间未检测到，可以考虑重置 smoothed_x/y/r 为 None
                if DEBUG_MODE >= 2:
                    print("当前帧未检测到符合条件的圆")


            # --- 绘制结果 (调试用) ---
            if target_found and smoothed_x is not None:
                # 绘制平滑后的圆
                img.draw_circle(int(smoothed_x), int(smoothed_y), int(smoothed_r), color=(0, 255, 0), thickness=2)
                img.draw_cross(int(smoothed_x), int(smoothed_y), color=(0, 255, 0), size=10, thickness=1)
                #img.draw_string(int(smoothed_x) + 10, int(smoothed_y) - 10, f"Target", color=(0,255,0))
            elif best_circle: # 如果当前帧找到了，但全局状态是false（比如刚启动），也画一下原始的
                 img.draw_circle(best_circle.x(), best_circle.y(), best_circle.r(), color=(255, 0, 0), thickness=1)
                 img.draw_cross(best_circle.x(), best_circle.y(), color=(255, 0, 0), size=8, thickness=1)




            # --- 数据发送 ---
            # 发送格式: 帧头(0xAA) + 模式(0x03) + X坐标(1 byte) + Y坐标(1 byte) + 帧尾(0xBB)
            # 注意：X, Y 坐标被限制在 0-255 范围内！

            send_buf = bytearray(5) # 新格式为 5 字节
            send_buf[0] = 0xAA     # 帧头
            send_buf[1] = 0x03     # 状态/模式：粗加工区检测

            if target_found and smoothed_x is not None:
                # 获取平滑后的坐标
                cx_float = smoothed_x
                cy_float = smoothed_y

                # --- 重要：将坐标转换为整数并限制在 0-255 范围内 ---
                # 如果原始坐标范围超过 255，这里会丢失精度或信息
                # 例如，如果图像是 320x240，大于 255 的坐标会被截断为 255
                cx_byte = max(0, min(int(cx_float), 255))
                cy_byte = max(0, min(int(cy_float), 255))

                send_buf[2] = cx_byte & 0xFF # X坐标 (0-255)
                send_buf[3] = cy_byte & 0xFF # Y坐标 (0-255)

            else:
                # 未找到目标，发送预设值 (例如 0, 0)
                # STM32 端需要知道 (0, 0) 可能表示未找到目标
                send_buf[2] = 0x00 # X坐标 = 0
                send_buf[3] = 0x00 # Y坐标 = 0

            send_buf[4] = 0xBB     # 帧尾

            # 发送数据
            uart.write(send_buf)

            #可选的调试打印
            if DEBUG_MODE >= 1:
                print(f"发送数据 (New Format): {['0x{:02X}'.format(b) for b in send_buf]}")



    # 待机或其他模式处理
    elif current_operation == OPERATION_MODE['STANDBY']:
        pyb.LED(2).off() # 确保绿色LED关闭
        pyb.LED(3).off() # 确保蓝色LED关闭
        # 可以添加心跳或其他逻辑
        pass # 待机模式下通常不做图像处理

    # 其他模式 (LINE_TRACKING, BUFFER_AREA, QR_CODE) 在这里添加各自的逻辑...


    # 显示帧率 (调试用)
    # print(f"FPS: {clock.fps():.1f}")

